package urand

import "math/rand"

// Weights 权重列表
type Weights interface {
	TotalLen() int
	TotalWeight() int
	GetWeight(int) int
}

// RandInWeights 计算权重，并返回数组下标
func RandInWeights(weightList Weights) int {
	chance := rand.Intn(weightList.TotalWeight())
	var up int
	for i := 0; i < weightList.TotalLen(); i++ {
		up += weightList.GetWeight(i)
		if chance < up {
			return i
		}
	}
	return weightList.TotalLen() - 1
}

// BaseWeights 将 []int 包装为 Weights 接口
type BaseWeights []int

// TotalLen 总长度
func (b BaseWeights) TotalLen() int { return len(b) }

// TotalWeight 总权重
func (b BaseWeights) TotalWeight() int {
	total := 0
	for _, w := range b {
		total += w
	}
	return total
}

// GetWeight 获取指定权重
func (b BaseWeights) GetWeight(i int) int { return b[i] }

// Wrap 将 []int 包装为baseWeights
func Wrap(v []int) BaseWeights { return BaseWeights(v) }
